
#include "glodef.h"
#include "errlog.h"
#include "vecs.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include <petsc.h>
#include <petscvec.h>

#include "vecs_petsc.h"

#define CHKERR_B(n) if (n) {n = PetscError(__LINE__-1,__FUNCT__,__FILE__,__SDIR__,n,0," ");error_msg("%s: petsc error detected", __func__);}

int get_N_vec(vector A)
{
    return A->Size;
}

void check_use_vector(vector A)
{
    int low, high;
    VecGetOwnershipRange(A->VecXYZ, &low, &high);
    if((low!=0)||(high!=3*A->Size))
        error_msg("%s failed.\n", __func__);
}

double *get_XYZ_vec(vector A, double **X, double **Y, double **Z)
{
    int ierr;
    if(A->GetXYZtimes==0)
    {
        ierr=VecGetArray(A->VecXYZ, &A->ArrayXYZ);
        CHKERR_B(ierr);
    }
    A->GetXYZtimes++;
    if(X) *X=A->ArrayXYZ;
    if(Y) *Y=A->ArrayXYZ+A->Size;
    if(Z) *Z=A->ArrayXYZ+2*A->Size;
    return A->ArrayXYZ;
}

void restore_XYZ_vec(vector A, double **X, double **Y, double **Z)
{
    if((X==NULL)||(*X!=A->ArrayXYZ)||
            (Y==NULL)||(*Y!=A->ArrayXYZ+A->Size)||
            (Z==NULL)||(*Z!=A->ArrayXYZ+2*A->Size))
        error_msg("%s: must restore the same arrays.\n", __func__);
    A->GetXYZtimes--;
    if(A->GetXYZtimes==0)
    {
        int ierr=VecRestoreArray(A->VecXYZ, &A->ArrayXYZ);
        CHKERR_B(ierr);
    }
    if(A->GetXYZtimes<=0)
    {
        A->ArrayXYZ=NULL;
        A->GetXYZtimes=0;
    }
}

int get_N_scal(scalar A)
{
    return A->Size;
}

void check_use_scalar(scalar A)
{
    int low, high;
    VecGetOwnershipRange(A->VecF, &low, &high);
    if((low!=0)||(high!=A->Size))
        error_msg("%s failed.\n", __func__);
}
// void *get_Vec_scal(scalar A) {return A->VecF;}

double *get_F_scal(scalar A, double **F)
{
    if(A->GetFtimes==0)
    {
        int ierr=VecGetArray(A->VecF, &A->ArrayF);
        CHKERR_B(ierr);
    }
    A->GetFtimes++;
    if(F) *F=A->ArrayF;
    return A->ArrayF;
}

void restore_F_scal(scalar A, double **F)
{
    if(F==NULL|| *F!=A->ArrayF)
        error_msg("%s: must restore the same arrays.\n", __func__);
    A->GetFtimes--;
    if(A->GetFtimes==0)
    {
        int ierr=VecRestoreArray(A->VecF, &A->ArrayF);
        CHKERR_B(ierr);
    }
    if(A->GetFtimes<=0)
    {
        A->ArrayF=NULL;
        A->GetFtimes=0;
    }
}

// ==========================================================================

scalar vec2scal(vector Vec)
{
    if(Vec==NULL)
        return NULL;
    if(Vec->GetXYZtimes>0)
        error_msg("%s: must restore array first.\n", __func__);
    scalar retval=calloc(1, sizeof(sscalar));
    if(retval==NULL)
        error_msg("Memory allocation error in %s.\n", __func__);
    retval->Size=3*Vec->Size;
    retval->VecF=Vec->VecXYZ;
    retval->ArrayF=NULL;
    retval->GetFtimes=0;
    retval->owns_F=FALSE;
    return retval;
}

vector scal2vec(scalar Scal)
{
    if(Scal==NULL)
        return NULL;
    if(Scal->GetFtimes>0)
        error_msg("%s: must restore array first.\n", __func__);
    if(Scal->Size%3)
        error_msg("%s: cannot convert scalar with N=%d to vector.\n",
                  __func__, Scal->Size);
    vector retval=calloc(1, sizeof(sscalar));
    if(retval==NULL)
        error_msg("Memory allocation error in %s.\n", __func__);
    retval->Size=Scal->Size/3;
    retval->VecXYZ=Scal->VecF;
    retval->GetXYZtimes=0;
    retval->ArrayXYZ=NULL;
    retval->owns_XYZ=FALSE;
    return retval;
}

/****************************************************************************/

/*  utility fubctions */

inline static void print_petsc(Vec V, const char fname[])
{
    int ierr;
    PetscViewer file;
    ierr=PetscViewerASCIIOpen(PETSC_COMM_WORLD, fname, &file);
    CHKERR_B(ierr);
    ierr=VecView(V, file);
    CHKERR_B(ierr);
    ierr=PetscViewerDestroy(file);
    CHKERR_B(ierr);
}


/****************************************************************************/

/* allocate and release vector */

vector alloc_vec(int N)
{
    int ierr;
    if(N<=0) error_msg("Invalid data in %s.\n", __func__);

    vector A=calloc(1, sizeof(svector));
    if(A==NULL) error_msg("Memory allocation error in %s.\n", __func__);
    A->Size=N;
    A->ArrayXYZ=NULL;
    A->GetXYZtimes=0;
    ierr=VecCreate(PETSC_COMM_SELF, &A->VecXYZ);
    CHKERR_B(ierr);
    ierr=VecSetSizes(A->VecXYZ, PETSC_DECIDE, 3*N);
    CHKERR_B(ierr);
    ierr=VecSetType(A->VecXYZ, VECSEQ);
    CHKERR_B(ierr);
    A->owns_XYZ=TRUE;
    return A;
}

void free_vec(vector A)
{
    if(A==NULL)
        return;
    if(A->owns_XYZ)
    {
        int ierr=VecDestroy(A->VecXYZ);
        CHKERR_B(ierr);
    }
    memset(A, 0, sizeof(*A));
    free(A);
}

void copy_vec(const vector S, vector T) /* copies contents of S into T */
{
    assert((S!=NULL)&&(T!=NULL));
    assert(S->Size==T->Size);
    int ierr=VecCopy(S->VecXYZ, T->VecXYZ);
    CHKERR_B(ierr);
}

void clean_vec(vector A)
{
    assert(A!=NULL);
    int ierr=VecZeroEntries(A->VecXYZ);
    CHKERR_B(ierr);
}

void print_vec(vector A, char *fname)
{
    assert(A!=NULL);
    print_petsc(A->VecXYZ, fname);
}

/****************************************************************************/

/* vector operations */

void scale_vec(vector A, double alpha)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecScale(A->VecXYZ, alpha);
    CHKERR_B(ierr);
}

void axpy_vec(vector y, double a, vector x)
{
    if(y==NULL) error_msg("NULL pointer in %s.\n", __func__);
    if(x==NULL) return;
    if(x->Size!=y->Size) error_msg("Dimensions disagree in %s.\n", __func__);
    if(fabs(a)<SMALLF) return;
    int ierr=VecAXPY(y->VecXYZ, a, x->VecXYZ);
    CHKERR_B(ierr);
}

void waxpy_vec(vector w, double a, vector x, vector y)
{
    if(w==NULL||y==NULL||x==NULL)
        error_msg("NULL pointer in %s.\n", __func__);
    if(x->Size!=y->Size||x->Size!=w->Size)
        error_msg("Dimensions disagree in %s.\n", __func__);
    int ierr=VecWAXPY(w->VecXYZ, a, x->VecXYZ, y->VecXYZ);
    CHKERR_B(ierr);
}

double dot_vec(vector x, vector y)
{
    if(x==NULL||y==NULL) error_msg("NULL pointer in %s.\n", __func__);
    if(x->Size!=y->Size)
        error_msg("Dimensions disagree in %s.\n", __func__);
    double dot=0.0;
    int ierr=VecDot(x->VecXYZ, y->VecXYZ, &dot);
    CHKERR_B(ierr);
    return dot;
}

void set_value_vec(vector A, int pos, double *val)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecSetValue(A->VecXYZ, pos, val[0], INSERT_VALUES);
    CHKERR_B(ierr);
    pos+=A->Size;
    ierr=VecSetValue(A->VecXYZ, pos, val[1], INSERT_VALUES);
    CHKERR_B(ierr);
    pos+=A->Size;
    ierr=VecSetValue(A->VecXYZ, pos, val[2], INSERT_VALUES);
    CHKERR_B(ierr);
}

/****************************************************************************/

/* allocate and release scalar */

scalar alloc_scal(int N)
{
    if(N<=0) error_msg("Invalid data in %s.\n", __func__);

    scalar A=calloc(1, sizeof(sscalar));
    if(A==NULL) error_msg("Memory allocation error in %s.\n", __func__);
    A->Size=N;
    A->ArrayF=NULL;
    A->GetFtimes=0;
    int ierr=VecCreate(PETSC_COMM_SELF, &A->VecF);
    CHKERR_B(ierr);
    ierr=VecSetSizes(A->VecF, PETSC_DECIDE, N);
    CHKERR_B(ierr);
    ierr=VecSetType(A->VecF, VECSEQ);
    CHKERR_B(ierr);
    A->owns_F=TRUE;
    return A;
}

void free_scal(scalar A)
{
    if(A==NULL)
        return;
    if(A->owns_F)
    {
        int ierr=VecDestroy(A->VecF);
        CHKERR_B(ierr);
    }
    memset(A, 0, sizeof(*A));
    free(A);
}

void copy_scal(const scalar S, scalar T) /* copies contents of S into T */
{
    assert((S!=NULL)&&(T!=NULL));
    assert(S->Size==T->Size);
    int ierr=VecCopy(S->VecF, T->VecF);
    CHKERR_B(ierr);
}

void clean_scal(scalar A)
{
    assert(A!=NULL);
    int ierr=VecZeroEntries(A->VecF);
    CHKERR_B(ierr);
}

void print_scal(scalar A, char *fname)
{
    assert(A!=NULL);
    print_petsc(A->VecF, fname);
}

/****************************************************************************/

/* scalar operations */

void scale_scal(scalar A, double a)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecScale(A->VecF, a);
    CHKERR_B(ierr);
}

void axpy_scal(scalar y, double a, scalar x)
{
    if((x==NULL)||(y==NULL))
        error_msg("NULL pointer in %s.\n", __func__);
    if(x->Size!=y->Size)
        error_msg("Dimensions disagree in %s.\n", __func__);
    int ierr=VecAXPY(y->VecF, a, x->VecF);
    CHKERR_B(ierr);
}

void waxpy_scal(scalar w, double a, scalar x, scalar y)
{
    if((w==NULL)||(y==NULL)||(x==NULL))
        error_msg("NULL pointer in %s.\n", __func__);
    if((x->Size!=y->Size)||(x->Size!=w->Size))
        error_msg("Dimensions disagree in %s.\n", __func__);
    int ierr=VecWAXPY(w->VecF, a, x->VecF, y->VecF);
    CHKERR_B(ierr);
}

double dot_scal(scalar x, scalar y)
{
    if(x==NULL||y==NULL)
        error_msg("NULL pointer in %s.\n", __func__);
    if(x->Size!=y->Size)
        error_msg("Dimensions disagree in %s.\n", __func__);
    double dot=0.0;
    int ierr=VecDot(x->VecF, y->VecF, &dot);
    CHKERR_B(ierr);
    return dot;
}

double sum_scal(scalar x)
{
    if(x==NULL) error_msg("NULL pointer in %s.\n", __func__);
    double sum=0.0;
    int ierr=VecSum(x->VecF, &sum);
    CHKERR_B(ierr);
    return sum;
}

void set_value_scal(scalar A, int pos, double val)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecSetValue(A->VecF, pos, val, INSERT_VALUES);
    CHKERR_B(ierr);
}

void assemble_scal(scalar A)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecAssemblyBegin(A->VecF);
    CHKERR_B(ierr);
    ierr=VecAssemblyEnd(A->VecF);
    CHKERR_B(ierr);
}

void get_ownership_range_scal(scalar A, int *low, int *high)
{
    if(A==NULL) error_msg("NULL pointer in %s.\n", __func__);
    int ierr=VecGetOwnershipRange(A->VecF, low, high);
    CHKERR_B(ierr);
}

